package com.johnguo.beam_test.java;

import java.util.Arrays;
import java.util.List;

import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;

public class FlattenDemo {
	static final List<String> LINES1 = Arrays.asList("bye anna");

	static final List<String> LINES2 = Arrays.asList("hello john");

	public static void main(String... args) {
		PipelineOptions options = PipelineOptionsFactory.fromArgs(args).withValidation().create();
		Pipeline p = Pipeline.create(options);
		PCollection<String> reader1 = p.apply(Create.of(LINES1)).setCoder(StringUtf8Coder.of());
		PCollection<String> reader2 = p.apply(Create.of(LINES2)).setCoder(StringUtf8Coder.of());
		PCollectionList<String> collections = PCollectionList.of(reader1).and(reader2);
		PCollection<String> merged = collections.apply(Flatten.<String>pCollections());
		merged.apply(ParDo.of(new DoFn<String, KV<String, Long>>() {
			private static final long serialVersionUID = 1L;

			@ProcessElement
			public void processElement(ProcessContext c) {
				for (String word : c.element().split(" ")) {
					if (!word.isEmpty()) {
						System.out.println(word);
						c.output(KV.of(word, 1l));
					}
				}
			}
		}));
		p.run().waitUntilFinish();
	}
}
